python train_1p_with_amp.py --learning_rate 0.0001 --batch_size 1 --epochs 1